import os

import time
import re
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from utils.title_and_table_change import title_and_table_change
import logging
from log_config import setup_logging
from utils.table_merge.table_md_merge import remove_body_html_tags

# 设置日志
setup_logging()

 
label = [1,2,3,4,6]    
    

for i in label:
# for i in [18]:
    start_time = time.time()
    # 参数设置:文件路径、文件名

    # pdf_file_name = f"feature_table/{i}.pdf"  # 替换真实的pdf路径
    pdf_file_name = f"input_doc/sample/提取自26.pdf"  # 替换真实的pdf路径
    logging.info(f'运行{pdf_file_name:}')
    # pdf_file_name = "/home/fengjie/MinerU/demo/pdf_files/剪切版-非跨页.pdf"  # 替换真实的pdf路径

    name_without_suff = pdf_file_name.split(".")[0]

    # 准备环境
    local_image_dir, local_md_dir = "output/images", "output"
    image_dir = str(os.path.basename(local_image_dir)) # image_dir = images
    ## 创建存储文件夹
    os.makedirs(local_image_dir, exist_ok=True)
    ## 实例化writer对象：用于把数据写入文件中
    image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(
        local_md_dir
    )

    # read bytes
    reader1 = FileBasedDataReader("")
    pdf_bytes = reader1.read(pdf_file_name)  # 准备pdf目录

    # 处理过程
    ## 创造数据集实例
    ds = PymuDocDataset(pdf_bytes)


    # 模型推理
    infer_result = ds.apply(doc_analyze, ocr=True)
    # 管道处理
    pipe_result = infer_result.pipe_ocr_mode(image_writer)

    # 保持处理前数据
    md_content = pipe_result.get_markdown(image_dir)
    md_content = remove_body_html_tags(md_content)
    md_writer.write_string(f"{name_without_suff}_pre.md",md_content)

    # md_content = title_and_table_change(pipe_result,local_image_dir,image_dir)
    try:
        # 执行目标代码
        md_content = title_and_table_change(pipe_result, local_image_dir, image_dir)
        logging.info("title_and_table_change executed successfully.")
    except Exception as e:
        # 捕获异常并记录日志
        logging.error(f"Error occurred while executing title_and_table_change: {e}", exc_info=True)
        md_content = '程序错误，转换失败！！！'
    
    # 保存到markdown文件
    md_writer.write_string(f"{name_without_suff}.md",md_content)
    # pipe_result.draw_layout(os.path.join(local_md_dir, f"{name_without_suff}_layout.pdf"))

    # ### dump middle json
    # pipe_result.dump_middle_json(md_writer, f'{name_without_suff}_middle.json')

    end_time = time.time()
    logging.info(f"time:{end_time-start_time}")
    break

